from ext import limiter
from flask import request,current_app
from base.authentication import login_required

'1、限制的方法，响应什么就按照什么进行限制'
#获取请求的IP地址，设置成默认的全局流量控制
def get_remote_ip():
    real_ip = request.headers.get('X-Real-IP')
    forwarded_for = request.headers.get('X-Forwarded-For')
    local_ip = request.remote_addr
    if forwarded_for:
        ip = forwarded_for.split(',')[0]
    elif real_ip:
        ip = real_ip
    else:
        ip = local_ip
    return ip

#根据用户的id进行限制
@login_required
def get_user_id():
    token = request.cookies.get('token')
    if not token:
        token = request.headers.get('token')
    print('访问的用户',request.user)
    try:
        user = request.user
        if user:
            return user.id
        else:
            return token
    except Exception as _:
        return token

'2、流量控制的装饰器函数'
#用户访问的流量控制: 需要用户已经登录
user_limiter = limiter.limit('10/minute',key_func=get_user_id)
#注册和登录接口的流量控制：无需用户登录
login_limiter = limiter.limit('5/minute',key_func=get_remote_ip)
register_limiter = limiter.limit('5/minute',key_func=get_remote_ip)